from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from app.core.config import settings
from app.core.logger import attach_sql_logger, logger

# 创建异步数据库引擎
_engine = create_async_engine(
    settings.DATABASE_URL,
    echo=False,
    pool_pre_ping=True,  # 开启连接预检查
    pool_recycle=3600,  # 60分钟回收连接
)
attach_sql_logger(_engine)

# 创建异步会话工厂
AsyncSessionLocal = sessionmaker(
    bind=_engine,
    class_=AsyncSession,
    expire_on_commit=False,
)

# 上下午管理器：任务专用会话
class TaskDBSession:
    """
    任务模块数据库会话管理器
    """
    def __init__(self):
        self.session = None

    async def __aenter__(self):
        self.session = AsyncSessionLocal()
        logger.info("Opened new task DB session")
        return self.session

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if exc_val:
            await self.session.rollback()
            logger.error(f"Task DB session rollback: {exc_val}")
        else:
            await self.session.commit()
        await self.session.close()
        logger.info("Closed task DB session")

async def get_task_db():
    async with TaskDBSession() as session:
        yield session